package com.xiashitech.agent.instrumentation.db.resultset.inst;

import com.xiashitech.agent.cache.status.CacheStatus;
import com.xiashitech.agent.constants.Instrumentation;
import com.xiashitech.agent.instrumentation.db.common.XSResultSizeContext;
import com.xiashitech.agent.instrumentation.helper.calldepth.XSResultSetStatementCallDepth;
import com.xiashitech.agent.instrumentation.utils.BaggageUtil;
import com.xiashitech.agent.utils.DurationPrinter;
import io.opentelemetry.javaagent.bootstrap.CallDepth;
import io.opentelemetry.javaagent.bootstrap.internal.InstrumentationConfig;
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
import io.opentelemetry.javaagent.shaded.instrumentation.jdbc.internal.JdbcData;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Context;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;

import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.Statement;

import static com.xiashitech.agent.instrumentation.db.resultset.XSResultSetSingletons.instrumenter;
import static io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge.currentContext;
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.implementsInterface;
import static net.bytebuddy.matcher.ElementMatchers.*;

public class XSResultSetStatementInstrumentation implements TypeInstrumentation {

    @Override
    public ElementMatcher<ClassLoader> classLoaderOptimization() {
        return hasClassesNamed("java.sql.Statement","java.sql.PreparedStatement");
    }

    @Override
    public ElementMatcher<TypeDescription> typeMatcher() {
        return implementsInterface(named("java.sql.Statement")).or(implementsInterface(named("java.sql.PreparedStatement")));
    }

    @Override
    public void transform(TypeTransformer transformer) {
        transformer.applyAdviceToMethod(
                named("executeQuery")
                        .and(takesArgument(0, String.class))
                        .and(takesArguments(1))
                        .and(returns(named("java.sql.ResultSet")))
                        .and(isPublic()),
                XSResultSetStatementInstrumentation.class.getName() + "$XSStatementResultSetAdvice");
        transformer.applyAdviceToMethod(
                named("getResultSet")
                        .and(takesNoArguments())
                        .and(returns(named("java.sql.ResultSet")))
                        .and(isPublic()),
                XSResultSetStatementInstrumentation.class.getName() + "$XSResultSetAdvice");
        transformer.applyAdviceToMethod(
                named("executeQuery")
                        .and(takesNoArguments())
                        .and(returns(named("java.sql.ResultSet")))
                        .and(isPublic()),
                XSResultSetStatementInstrumentation.class.getName() + "$XSResultSetAdvice");
    }

    @SuppressWarnings("unused")
    public static class XSStatementResultSetAdvice {

        @Advice.OnMethodEnter(suppress = Throwable.class)
        public static void onEnterForResultSetPrepare(
                @Advice.This Statement statement,
                @Advice.Argument(0) String sql,
            @Advice.Local("otelCallDepth") CallDepth callDepth) {
//            if(statement != null && statement.getClass().getName().startsWith("xiashitech.shade.extload")) {
//                return;
//            }
            DurationPrinter durationPrinter = new DurationPrinter("XSResultSetStatementInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            String suppressSpecialsql = InstrumentationConfig.get().getString("xs.suppress.specialsql", Instrumentation.DEFAULT_SUPPRESS_SQL);
            if(!suppressSpecialsql.contains(Instrumentation.DEFAULT_SUPPRESS_SQL)) {
                suppressSpecialsql += ","+Instrumentation.DEFAULT_SUPPRESS_SQL;
            }
            if(sql != null && suppressSpecialsql.contains(sql)) {
                return;
            }
            callDepth = CallDepth.forClass(XSResultSetStatementCallDepth.class);
            callDepth.getAndIncrement();
            durationPrinter.printElapseDuration("XSStatementResultSetAdvice b");
        }

        /**
         * 记录结果集查询条数的开始入口
         * @param throwable
         * @param resultSet
         */
        @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
        public static void onExitForResultSetStart(
                @Advice.This Statement statement,
                @Advice.Argument(0) String sql,
                @Advice.Thrown Throwable throwable,
                @Advice.Return ResultSet resultSet,
                @Advice.Local("otelCallDepth") CallDepth callDepth) {
//            if(statement != null && statement.getClass().getName().startsWith("xiashitech.shade.extload")) {
//                return;
//            }
            DurationPrinter durationPrinter = new DurationPrinter("XSResultSetStatementInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            String suppressSpecialsql = InstrumentationConfig.get().getString("xs.suppress.specialsql", Instrumentation.DEFAULT_SUPPRESS_SQL);
            if(!suppressSpecialsql.contains(Instrumentation.DEFAULT_SUPPRESS_SQL)) {
                suppressSpecialsql += ","+Instrumentation.DEFAULT_SUPPRESS_SQL;
            }
            if(sql != null && suppressSpecialsql.contains(sql)) {
                return;
            }
            if(callDepth.decrementAndGet() > 0) {
                return;
            }
            if(XSResultSizeContext.getResultSetContext(resultSet) != null) {
                return;
            }
            XSResultSizeContext.Request request = XSResultSizeContext.Request.builder().resultSet(resultSet).size(0L).build();
            Context parentContext = currentContext();
            if (!instrumenter().shouldStart(parentContext, request)) {
                return;
            }
            parentContext = BaggageUtil.syncBaggage().build().storeInContext(parentContext);
            Context context = instrumenter().start(parentContext, request);
            context = BaggageUtil.syncBaggage().build().storeInContext(context);
            XSResultSizeContext.addResultSetContext(XSResultSizeContext.ResultSetContext.builder().resultSet(resultSet).context(context).throwable(throwable).build());
            durationPrinter.printElapseDuration("XSStatementResultSetAdvice e");
        }
    }

    @SuppressWarnings("unused")
    public static class XSResultSetAdvice {

        @Advice.OnMethodEnter(suppress = Throwable.class)
        public static void onEnterForResultSetPrepare(
                @Advice.This Statement statement,
                @Advice.Local("otelCallDepth") CallDepth callDepth) {
//            if(statement != null && statement.getClass().getName().startsWith("xiashitech.shade.extload")) {
//                return;
//            }
            DurationPrinter durationPrinter = new DurationPrinter("XSResultSetStatementInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            callDepth = CallDepth.forClass(XSResultSetStatementCallDepth.class);
            callDepth.getAndIncrement();
            String sql = null;
            if (statement instanceof PreparedStatement) {
                sql = JdbcData.preparedStatement.get((PreparedStatement)statement);
            } else if(statement instanceof Statement) {
                XSResultSizeContext.StatementSQLContext statementSQLContext = XSResultSizeContext.getStatementSQLContext(statement);
                if(statementSQLContext != null) {
                    sql = statementSQLContext.getSql();
                }
            }
            String suppressSpecialsql = InstrumentationConfig.get().getString("xs.suppress.specialsql", Instrumentation.DEFAULT_SUPPRESS_SQL);
            if(!suppressSpecialsql.contains(Instrumentation.DEFAULT_SUPPRESS_SQL)) {
                suppressSpecialsql += ","+Instrumentation.DEFAULT_SUPPRESS_SQL;
            }
            durationPrinter.printElapseDuration("XSResultSetAdvice b");
            if(sql != null && suppressSpecialsql.contains(sql)) {
                return;
            }
        }

        /**
         * 记录结果集查询条数的开始入口
         * @param throwable
         * @param resultSet
         */
        @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
        public static void onExitForResultSetStart(
                @Advice.This Statement statement,
                @Advice.Thrown Throwable throwable,
                @Advice.Return ResultSet resultSet,
                @Advice.Local("otelCallDepth") CallDepth callDepth) {
//            if(statement != null && statement.getClass().getName().startsWith("xiashitech.shade.extload")) {
//                return;
//            }
            DurationPrinter durationPrinter = new DurationPrinter("XSResultSetStatementInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            if(callDepth.decrementAndGet() > 0) {
                return;
            }
            String sql = null;
            if (statement instanceof PreparedStatement) {
                sql = JdbcData.preparedStatement.get((PreparedStatement)statement);
            } else if(statement instanceof Statement) {
                XSResultSizeContext.StatementSQLContext statementSQLContext = XSResultSizeContext.getStatementSQLContext(statement);
                if(statementSQLContext != null) {
                    sql = statementSQLContext.getSql();
                }
            }
            String suppressSpecialsql = InstrumentationConfig.get().getString("xs.suppress.specialsql", Instrumentation.DEFAULT_SUPPRESS_SQL);
            if(!suppressSpecialsql.contains(Instrumentation.DEFAULT_SUPPRESS_SQL)) {
                suppressSpecialsql += ","+Instrumentation.DEFAULT_SUPPRESS_SQL;
            }
            if(sql != null && suppressSpecialsql.contains(sql)) {
                return;
            }
            if(XSResultSizeContext.getResultSetContext(resultSet) != null) {
                return;
            }
            XSResultSizeContext.Request request = XSResultSizeContext.Request.builder().resultSet(resultSet).size(0L).build();
            Context parentContext = currentContext();
            if (!instrumenter().shouldStart(parentContext, request)) {
                return;
            }
            parentContext = BaggageUtil.syncBaggage().build().storeInContext(parentContext);
            Context context = instrumenter().start(parentContext, request);
            context = BaggageUtil.syncBaggage().build().storeInContext(context);
            XSResultSizeContext.addResultSetContext(XSResultSizeContext.ResultSetContext.builder().resultSet(resultSet).context(context).throwable(throwable).build());
            durationPrinter.printElapseDuration("XSResultSetAdvice e");
        }
    }
}
